"""
由可行域体素转节点图
@author: CeasarSmj
@date: 2025-4-16

算法流程；
1. 读取 ply 文件, 读取里面的 voxel grid 数据。
2. 初始化节点图 graph = {"nodes": [], edges: []},
    其中 nodes 的元素 node = {"id": int, "x": float, "y": float, "z": float}
    edges 的元素 edge = {"id": int, "from": int, "to": int}
3. 随机挑选1个体素放入列表 block_list[i] 中
4. 计算  block_list[i] 中所有点的均值中心,然后持续添加和中心距离小于 block_size/2 且连通的体素, 直到不能再添加
5. 检查是否循环了3次, 如果有则执行 6 , 否则回到 4
6. 检查 voxel_grid 中是否还有 voxel, 如果有则执行 3, 否则执行 7
7. 计算所有 block_list 中点的均值中心, 作为节点添加到 graph 中。
8. 计算 block_list 中每两块block之间的连通性, 如果连通则作为边添加到 graph 中。
9. 把 graph 用 yaml 格式保存到 output_file 中。

注1: ply 文件有一个和它同名的 yaml 文件, 里面有分辨率和坐标系参数
    label = ['resolution','x','y','z','roll','pitch','yaw'。
注2: 可行域地图中可行的 voxel 值为 1 , 所以删除等价于赋值为 0
注3: 这里连通的定义为: 距离小于 max_connect_dist的体素视为相邻, 连通是相邻的传递闭包
"""

import os
import yaml
import numpy as np
from scipy.spatial import KDTree
import open3d as o3d


class Voxel2graph:
    def __init__(self, block_size=4.0, max_connect_dist=0.3):
        self.block_size = block_size
        self.max_connect_dist = max_connect_dist
        self.graph = {"nodes": [], "edges": []}
        self.voxel_grid = None
        self.config = None

    def __call__(self, input_dir, output_file, if_visualize=False):
        print("Loading config...")
        self.config = self.load_config(input_dir)
        print("Loading voxel grid...")
        self.voxel_grid = self.load_voxel_grid(input_dir)
        print("Generating graph...")
        self.generate_graph()
        print("Saving graph...")
        self.save_graph(output_file)
        if if_visualize:
            print("Visualizing graph...")
            self.visualize_graph()
        print("Done!")

    def load_voxel_grid(self, file_path):
        pcd = o3d.io.read_point_cloud(file_path)
        points = np.asarray(pcd.points)
        resolution = self.config.get("resolution", 0.0)
        x0 = self.config.get("x", 0.0)
        y0 = self.config.get("y", 0.0)
        z0 = self.config.get("z", 0.0)
        actual_points = points * resolution + np.array([x0, y0, z0])
        return actual_points

    def load_config(self, file_path):
        config_path = file_path.replace(".ply", ".yaml")
        config = {
            "resolution": 0.0,
            "x": 0.0,
            "y": 0.0,
            "z": 0.0,
            "roll": 0.0,
            "pitch": 0.0,
            "yaw": 0.0,
        }
        if os.path.exists(config_path):
            with open(config_path, "r") as f:
                loaded_config = yaml.safe_load(f)
                config.update(loaded_config)
        return config

    def generate_graph(self):
        voxel_grid = self.voxel_grid
        if voxel_grid.size == 0:
            return
        n_voxels = voxel_grid.shape[0]
        mask = np.zeros(n_voxels, dtype=bool)
        block_list = []

        while not np.all(mask):
            remaining_indices = np.where(~mask)[0]
            if not remaining_indices.size:
                break
            seed_idx = np.random.choice(remaining_indices)
            current_block = [voxel_grid[seed_idx]]
            mask[seed_idx] = True

            for _ in range(3):
                added = True
                while added:
                    added = False
                    center = np.mean(current_block, axis=0)
                    remaining_indices = np.where(~mask)[0]
                    if not remaining_indices.size:
                        break
                    remaining_voxels = voxel_grid[remaining_indices]

                    dist_to_center = np.linalg.norm(remaining_voxels - center, axis=1)
                    in_radius_mask = dist_to_center <= (self.block_size / 2)
                    in_radius_voxels = remaining_voxels[in_radius_mask]
                    in_radius_indices = remaining_indices[in_radius_mask]

                    if not in_radius_voxels.size:
                        break

                    current_block_arr = np.array(current_block)
                    tree = KDTree(current_block_arr)
                    dists, _ = tree.query(in_radius_voxels, k=1)
                    connected_mask = dists <= self.max_connect_dist
                    connected_voxels = in_radius_voxels[connected_mask]
                    connected_indices = in_radius_indices[connected_mask]

                    if connected_voxels.size > 0:
                        current_block.extend(connected_voxels.tolist())
                        mask[connected_indices] = True
                        added = True

            block_list.append(np.array(current_block))
            print(f"Processing block {len(block_list)}")

        nodes = []
        for i, block in enumerate(block_list):
            center = np.mean(block, axis=0)
            nodes.append(
                {
                    "id": i,
                    "x": float(center[0]),
                    "y": float(center[1]),
                    "z": float(center[2]),
                }
            )

        edges = []
        edge_id = 0
        block_trees = [KDTree(block) for block in block_list]

        is_isolated = [True for _ in range(len(block_list))]
        for i in range(len(block_list)):
            for j in range(i + 1, len(block_list)):
                distances, _ = block_trees[i].query(block_list[j], k=1)
                if np.any(distances <= self.max_connect_dist):
                    edges.append(
                        {"id": edge_id, "from": nodes[i]["id"], "to": nodes[j]["id"]}
                    )
                    edge_id += 1
                    is_isolated[i] = False
                    is_isolated[j] = False

        nodes = [nodes[i] for i in range(len(nodes)) if is_isolated[i] is not True]

        self.graph["nodes"] = nodes
        self.graph["edges"] = edges

    def save_graph(self, output_file):
        graph = self.graph_transform(self.graph)
        with open(output_file, "w") as f:
            yaml.dump(graph, f, default_flow_style=False)

    def graph_transform(self, graphin):
        id2index = {node["id"]: i for i, node in enumerate(graphin["nodes"])}

        nodes = [[node["x"], node["y"], node["z"]] for node in graphin["nodes"]]

        edges = [[id2index[edge['from']], id2index[edge['to']]] for edge in graphin["edges"]]

        graphout = {"nodes": nodes, "edges": edges}
        return graphout

    def visualize_graph(self):
        vis = o3d.visualization.Visualizer()
        vis.create_window()
        nodes = self.graph["nodes"]
        if nodes:
            points = np.array([[n["x"], n["y"], n["z"]] for n in nodes])
            pcd = o3d.geometry.PointCloud()
            pcd.points = o3d.utility.Vector3dVector(points)
            pcd.paint_uniform_color([1, 0, 0])  # Red for nodes
            vis.add_geometry(pcd)

            edges = self.graph["edges"]
            if edges:
                line_points = []
                lines = []
                for idx, edge in enumerate(edges):
                    from_node = [node for node in nodes if node["id"] == edge["from"]][
                        0
                    ]
                    to_node = [node for node in nodes if node["id"] == edge["to"]][0]
                    line_points.append([from_node["x"], from_node["y"], from_node["z"]])
                    line_points.append([to_node["x"], to_node["y"], to_node["z"]])
                    lines.append([2 * idx, 2 * idx + 1])

                line_set = o3d.geometry.LineSet()
                line_set.points = o3d.utility.Vector3dVector(line_points)
                line_set.lines = o3d.utility.Vector2iVector(lines)
                line_set.paint_uniform_color([0, 0, 1])  # Blue for edges
                vis.add_geometry(line_set)

        vis.run()
        vis.destroy_window()


if __name__ == "__main__":
    voxel_path = "voxel_maps/map.ply"
    graph_path = "node_maps/graph.yaml"

    voxel2graph = Voxel2graph()
    voxel2graph(voxel_path, graph_path, if_visualize=True)
